
package com.example.demo.demos.rest;

import java.net.http.HttpClient;
import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.HashMap;
import java.util.Map;

import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpAsyncClient;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.spec.McpSchema;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.mcp.client.autoconfigure.NamedClientMcpTransport;
import org.springframework.ai.mcp.client.autoconfigure.SseHttpClientTransportAutoConfiguration;
import org.springframework.ai.mcp.client.autoconfigure.SseWebFluxTransportAutoConfiguration;
import org.springframework.ai.mcp.client.autoconfigure.StdioTransportAutoConfiguration;
import org.springframework.ai.mcp.client.autoconfigure.configurer.McpAsyncClientConfigurer;
import org.springframework.ai.mcp.client.autoconfigure.properties.McpClientCommonProperties;
import org.springframework.ai.mcp.client.autoconfigure.properties.McpSseClientProperties;
import org.springframework.ai.mcp.customizer.McpAsyncClientCustomizer;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.client.loadbalancer.LoadBalancerRequestFactory;
import org.springframework.context.annotation.Bean;

@AutoConfiguration(after = {StdioTransportAutoConfiguration.class, SseHttpClientTransportAutoConfiguration.class,
        SseWebFluxTransportAutoConfiguration.class})
@ConditionalOnClass({McpSchema.class})
@EnableConfigurationProperties(McpClientCommonProperties.class)
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true",
        matchIfMissing = true)
public class McpClientDetectAutoConfig {
    private Logger LOG = LoggerFactory.getLogger(McpClientDetectAutoConfig.class);

    private static String connectedClientName(String clientName, String serverConnectionName) {
        return clientName + " - " + serverConnectionName;
    }

    @Bean
    @ConditionalOnMissingBean
    @ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
    McpAsyncClientConfigurer mcpAsyncClientConfigurer(ObjectProvider<McpAsyncClientCustomizer> customizerProvider) {
        return new McpAsyncClientConfigurer(customizerProvider.orderedStream().toList());
    }
    // Async client configuration

    @Bean
    @ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
    public McpAsyncClientsWrapper mcpAsyncClientsWrapper(
            McpAsyncClientConfigurer mcpAsyncClientConfigurer,
            McpClientCommonProperties commonProperties,
            McpSseClientProperties sseProperties,
            ObjectProvider<ObjectMapper> objectMapperProvider,
            LoadBalancerClient discoveryClient
    ) {
        Map<String, McpAsyncClientWrapper> mccss = new HashMap<>();
        ObjectMapper o = objectMapperProvider.getIfAvailable(ObjectMapper::new);

        for (Map.Entry<String, McpSseClientProperties.SseParameters> sp : sseProperties.getConnections().entrySet()) {
            mccss.put(sp.getKey()
                    , new McpAsyncClientWrapper(
                            createMcpClient(sp.getKey(), sp.getValue(), o, commonProperties, mcpAsyncClientConfigurer, discoveryClient)
                            , sp.getKey(), sp.getValue(), o, commonProperties, mcpAsyncClientConfigurer,discoveryClient));
        }

        return new McpAsyncClientsWrapper(mccss);
    }


    @JsonInclude(JsonInclude.Include.NON_ABSENT)
    @JsonIgnoreProperties(ignoreUnknown = true)
    public record McpAsyncClientsWrapper(Map<String, McpAsyncClientWrapper> mccs) {
        public McpAsyncClient get(String serverid) {
            if (mccs.containsKey(serverid)) {
                return mccs.get(serverid).client();
            }
            return null;
        }
    }

    @JsonInclude(JsonInclude.Include.NON_ABSENT)
    @JsonIgnoreProperties(ignoreUnknown = true)
    public class McpAsyncClientWrapper {
        private McpAsyncClient client;
        private String serverid;
        private McpSseClientProperties.SseParameters config;
        private ObjectMapper objectMapper;
        private McpClientCommonProperties commonProperties;
        private McpAsyncClientConfigurer mcpAsyncClientConfigurer;
        private LoadBalancerClient discoveryClient;

        public McpAsyncClientWrapper(McpAsyncClient c, String s,
                                     McpSseClientProperties.SseParameters f
                , ObjectMapper o, McpClientCommonProperties com, McpAsyncClientConfigurer mc, LoadBalancerClient d) {
            client = c;
            serverid = s;
            config = f;
            objectMapper = o;
            commonProperties = com;
            mcpAsyncClientConfigurer = mc;
            discoveryClient = d;
        }

        private static Logger LOG = LoggerFactory.getLogger(McpAsyncClientWrapper.class);

        public McpAsyncClient client() {

            try {
                client.ping().block(Duration.of(20, ChronoUnit.MILLIS));
            }
            // 假设 hasFailed 方法用于检测客户端是否失败
            catch (Throwable x) {
                // 从 Spring 上下文中获取一个新的 McpAsyncClient 实例
                LOG.error("客户端client.ping", x);
                try {
                    client.close();
                } catch (final Exception e) {
                    LOG.error("关闭旧客户端失败", e);
                }
                client = reCreateMcpClient();
            }
            return client;
        }

        //        private set
        private synchronized McpAsyncClient reCreateMcpClient() {
            return createMcpClient(serverid, config, objectMapper, commonProperties, mcpAsyncClientConfigurer,discoveryClient);
            // 更新MCP客户端引用和工具回调
//        mcpClient = newClient;
//        toolCallbacks = List.of(new SyncMcpToolCallbackProvider(mcpClient).getToolCallbacks());
        }
    }

    public static McpAsyncClient createMcpClient(String s,
                                                 McpSseClientProperties.SseParameters f
            , ObjectMapper o, McpClientCommonProperties com, McpAsyncClientConfigurer mc, LoadBalancerClient dc) {
        String baseUrl = f.url();
        String sseEndpoint = f.sseEndpoint() != null
                ? f.sseEndpoint() : "/sse";
        var transport = HttpClientSseClientLoadBalanceTransport.builder(baseUrl)
                .sseEndpoint(sseEndpoint)
                .clientBuilder(HttpClient.newBuilder())
                .objectMapper(o)
                .dc(dc)
                .build();
        NamedClientMcpTransport namedTransport = new NamedClientMcpTransport(s, transport);
        McpSchema.Implementation clientInfo = new McpSchema.Implementation(
                connectedClientName(com.getName(), namedTransport.name()),
                com.getVersion());

        McpClient.AsyncSpec spec = McpClient.async(namedTransport.transport())
                .clientInfo(clientInfo)
                .requestTimeout(com.getRequestTimeout());

        spec = mc.configure(namedTransport.name(), spec);

        var client = spec.build();

        if (com.isInitialized()) {
            client.initialize().block();
        }
        return client;
        // 更新MCP客户端引用和工具回调
//        mcpClient = newClient;
//        toolCallbacks = List.of(new SyncMcpToolCallbackProvider(mcpClient).getToolCallbacks());
    }
}
